"""
Author: Morphlng
Date: 2024-04-05 22:18:45
LastEditTime: 2024-12-16 23:03:50
LastEditors: Morphlng
Description: Lidar observation
FilePath: /DrivingGym/src/driving_gym/environment/agent/obs/lidar_obs.py
"""

from __future__ import annotations

import gymnasium as gym
import numpy as np

from driving_gym.data.data_provider import Snapshot
from driving_gym.environment.agent.obs.base_obs import BaseObs
from driving_gym.environment.scenario.base_scenario import BaseScenario
from driving_gym.misc.util import override
from driving_gym.simulation.adapter_interface import AdapterInterface


class LidarObs(BaseObs):
    def __init__(self, config: dict, adapter: AdapterInterface):
        super().__init__(config, adapter)

        self.source: str = config["source"]
        if self.source.startswith(self.actor_id + ":"):
            self.source = self.source[len(self.actor_id) + 1 :]

        self.lidar_bin: float = config.get("lidar_bin", 0.125)
        self.d_behind: float = config.get("d_behind", 16)
        self.obs_range: float = config.get("obs_range", 32)
        self.lidar_height: float = config.get("lidar_height", 2.1)
        self.normalize: bool = config.get("normalize", True)

        self.x_bins = np.arange(
            -(self.obs_range - self.d_behind),
            self.d_behind + self.lidar_bin,
            self.lidar_bin,
        )
        self.y_bins = np.arange(
            -self.obs_range / 2, self.obs_range / 2 + 0.01, self.lidar_bin
        )
        self.z_bins = [-self.lidar_height - 1, -self.lidar_height + 0.25, 1]

    @override(BaseObs)
    def get_obs(self, snapshot: Snapshot, scenario: BaseScenario) -> np.ndarray:
        lidar = snapshot.data[self.actor_id]["sensors"][self.source]

        lidar = lidar[:, :3]
        lidar, _ = np.histogramdd(lidar, bins=(self.x_bins, self.y_bins, self.z_bins))
        lidar[:, :, 0] = np.array(lidar[:, :, 0] > 0, dtype=np.uint8)
        lidar[:, :, 1] = np.array(lidar[:, :, 1] > 0, dtype=np.uint8)
        padding = np.zeros((lidar.shape[0], lidar.shape[1], 1), dtype=np.uint8)
        lidar = np.concatenate((lidar, padding), axis=2)
        lidar = np.flip(lidar, axis=1)
        lidar = np.rot90(lidar, 1)
        if self.normalize:
            lidar = lidar.astype(np.float32)
        else:
            lidar *= 255
            lidar = lidar.astype(np.uint8)
        return lidar

    @override(BaseObs)
    def get_observation_space(self) -> gym.spaces.Box:
        if self.normalize:
            low, high = 0, 1
            dtype = np.float32
        else:
            low, high = 0, 255
            dtype = np.uint8

        obs_size = int(self.obs_range / self.lidar_bin)
        shape = (obs_size, obs_size, 3)
        return gym.spaces.Box(low=low, high=high, shape=shape, dtype=dtype)
